import uuid
from typing import List
from typing import Literal
from typing import Annotated
from typing_extensions import TypedDict
from pydantic import BaseModel

from langchain_core.messages import SystemMessage
from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages

from utils import init_llm, set_env, visualize_graph

# 系统提示模板 - 用于指导AI从用户那里收集信息
template = """您的工作是从用户那里获取关于他们想要创建的提示模板类型的信息。

您应该从他们那里获取以下信息：

- 提示的目标是什么
- 哪些变量将传递给提示模板
- 输出不应该做什么的任何约束条件
- 输出必须遵守的任何要求

如果您无法辨别这些信息，请要求他们澄清！不要试图盲目猜测。

在您能够辨别所有信息后，调用相关工具。"""


class PromptInstructions(BaseModel):
    """提示大型语言模型的说明。
    定义了生成提示模板所需的四个关键要素：
    - objective: 提示的目标
    - variables: 变量列表
    - constraints: 约束条件
    - requirements: 要求列表
    """

    objective: str
    variables: List[str]
    constraints: List[str]
    requirements: List[str]


def info_chain(state):
    """
    信息收集节点 - 负责从用户那里收集生成提示模板所需的信息
    
    Args:
        state: 当前状态，包含消息历史
        
    Returns:
        包含AI响应的新状态
    """
    messages = [SystemMessage(content=template)] + state["messages"]
    response = llm_with_tool.invoke(messages)
    return {"messages": [response]}

# 获取提示消息的功能
# 只会在工具调用后收到消息。
def get_prompt_messages(messages: list):
    """
    从消息历史中提取工具调用参数，并构建用于生成提示模板的消息
    
    Args:
        messages: 消息历史列表
        
    Returns:
        用于生成提示模板的消息列表
    """
    tool_call = None
    other_msgs = []

    # 新系统提示 - 用于指导AI生成提示模板
    prompt_system = """根据以下要求，写一个好的提示模板：

    {reqs}"""

    for m in messages:
        if isinstance(m, AIMessage) and m.tool_calls:
            # 提取工具调用的参数
            tool_call = m.tool_calls[0]["args"]
        elif isinstance(m, ToolMessage):
            # 跳过工具消息
            continue
        elif tool_call is not None:
            # 收集工具调用后的其他消息
            other_msgs.append(m)
    return [SystemMessage(content=prompt_system.format(reqs=tool_call))] + other_msgs


def prompt_gen_chain(state):
    """
    提示生成节点 - 根据收集到的信息生成最终的提示模板
    
    Args:
        state: 当前状态，包含消息历史
        
    Returns:
        包含生成提示模板的新状态
    """
    messages = get_prompt_messages(state["messages"])
    response = llm.invoke(messages)
    return {"messages": [response]}


def get_state(state):
    """
    状态路由函数 - 决定工作流的下一步走向
    
    Args:
        state: 当前状态
        
    Returns:
        下一个节点的名称或END
    """
    messages = state["messages"]
    if isinstance(messages[-1], AIMessage) and messages[-1].tool_calls:
        # 如果最后一条消息是AI消息且包含工具调用，则添加工具消息
        return "add_tool_message"
    elif not isinstance(messages[-1], HumanMessage):
        # 如果最后一条消息不是用户消息，则结束工作流
        return END
    # 否则继续收集信息
    return "info"


class State(TypedDict):
    """状态类型定义 - 定义工作流状态的结构"""
    messages: Annotated[list, add_messages]


def build_graph():
    """
    构建LangGraph工作流
    
    Returns:
        编译后的图对象
    """
    # 创建内存保存器，用于保存对话状态
    memory = MemorySaver()
    # 创建状态图
    workflow = StateGraph(State)
    
    # 添加节点
    workflow.add_node("info", info_chain)  # 信息收集节点
    workflow.add_node("prompt", prompt_gen_chain)  # 提示生成节点

    @workflow.add_node
    def add_tool_message(state: State):
        """
        工具消息添加节点 - 在工具调用后添加工具响应消息
        
        Args:
            state: 当前状态
            
        Returns:
            包含工具消息的新状态
        """
        return {
            "messages": [
                ToolMessage(
                    content="Prompt generated!",  # 工具响应内容
                    tool_call_id=state["messages"][-1].tool_calls[0]["id"],  # 关联工具调用ID
                )
            ]
        }

    # 添加条件边 - 根据状态决定下一步
    workflow.add_conditional_edges("info", get_state, ["add_tool_message", "info", END])
    # 添加普通边
    workflow.add_edge("add_tool_message", "prompt")  # 工具消息 -> 提示生成
    workflow.add_edge("prompt", END)  # 提示生成 -> 结束
    workflow.add_edge(START, "info")  # 开始 -> 信息收集
    
    # 编译图
    graph = workflow.compile(checkpointer=memory)
    return graph


if __name__ == "__main__":
    # 主程序入口
    set_env()  # 设置环境变量
    llm = init_llm(temperature=0)  # 初始化LLM，温度设为0以获得确定性输出
    llm_with_tool = llm.bind_tools([PromptInstructions])  # 绑定工具到LLM
    graph = build_graph()  # 构建工作流图
    visualize_graph(graph)
    
    # 创建配置，包含唯一的线程ID
    config = {"configurable": {"thread_id": str(uuid.uuid4())}}
    
    # 主交互循环
    while True:
        # 获取用户输入
        user = input("User (q/Q to quit): ")
        # print(f"User (q/Q to quit): {user}")
        
        # 检查退出条件
        if user in {"q", "Q"}:
            print("AI: Byebye")
            break
            
        # 初始化输出变量
        output = None
        
        # 流式处理用户输入
        for output in graph.stream(
            {"messages": [HumanMessage(content=user)]}, config=config, stream_mode="updates"
        ):
            # 获取并打印最后一条消息
            last_message = next(iter(output.values()))["messages"][-1]
            last_message.pretty_print()

        # 检查是否完成提示生成
        if output and "prompt" in output:
            print("Done!")
